package com.ztom.v2;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author ZhangTao
 */
public class Code75TopKFrequentElement {

    public int[] topKFrequent(int[] nums, int k) {
        List<int[]> fnums = count(nums);
        int[] res = new int[k];

        // 采用降序
        k -= 1;
        int l = 0;
        int r = fnums.size() - 1;
        while (l <= r) {
            int[] mid = partition(fnums, l, r, fnums.get(l + (int) (Math.random() * (r - l + 1)))[1]);
            if (k < mid[0]) {
                r = mid[0] - 1;
            } else if (k > mid[1]) {
                l = mid[1] + 1;
            } else {
                // 收集答案
                for (int i = 0; i <= k; i++) {
                    res[i] = fnums.get(i)[0];
                }
                break;
            }
        }
        return res;
    }

    public int[] partition(List<int[]> nums, int start, int end, int num) {
        if (start >= end) {
            return new int[]{start, end};
        }
        // 注意这里要降序来排, 方便取前k个数
        int l = start - 1;
        int r = end + 1;
        int index = start;
        while (index < r) {
            if (nums.get(index)[1] > num) {
                Collections.swap(nums, index++, ++l);
            } else if (nums.get(index)[1] < num) {
                Collections.swap(nums, index, --r);
            } else {
                index++;
            }
        }
        return new int[]{++l, --r};
    }

    public List<int[]> count(int[] nums) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        return map.entrySet().stream()
                .map(entry -> new int[]{entry.getKey(), entry.getValue()})
                .collect(Collectors.toList());
    }


    public int[] topKFrequent1(int[] nums, int k) {
        List<int[]> fnums = count(nums);
        // 构建大根堆
        int heapSize = fnums.size();
        for (int i = (fnums.size() - 2 >> 1); i >=0; i--) {
            heapify(fnums, i, heapSize);
        }

        int[] res = new int[k];
        // 依次从大根堆弹出前 k 个数即可
        for (int i = 0; i < k; i++) {
            Collections.swap(fnums, 0, --heapSize);
            res[i] = fnums.get(heapSize)[0];
            heapify(fnums, 0, heapSize);
        }
        return res;
    }

    public void heapify(List<int[]> fnums, int index, int size) {
        int left = (index << 1) + 1;
        while (left < size) {
            int largest = left + 1 < size && fnums.get(left + 1)[1] > fnums.get(left)[1] ? left + 1 : left;
            largest = fnums.get(largest)[1] > fnums.get(index)[1] ? largest : index;
            if (largest == index) {
                break;
            }
            Collections.swap(fnums, largest, index);
            index = largest;
            left = (index << 1) + 1;
        }
    }
}
